from flask import Flask, jsonify, send_from_directory, request, render_template, Response
import os
import subprocess
import requests
import re
from pathlib import Path
import tempfile
import shutil
import json
import openai  # 需要先安装 openai 包: pip install openai
import configparser  # 添加 configparser 导入

app = Flask(__name__)

# 设置目录
FILES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "converted_files")

# 添加配置加载和保存功能
# 配置文件名
CONFIG_FILE = 'config.properties'

# 默认配置
DEFAULT_CONFIG = {
    'API': {
        'base_url': 'https://api.farai.com/v1',
        'dataset_id': 'default_dataset',
        'document_id': 'default_document',
        'api_key': 'your_api_key_here'
    },
    'OpenAI': {
    },
    'Settings': {
        'upload.dir': './converted_files',
        'allowed.extensions': '.md,.pdf,.txt,.doc,.docx',
        'output_format': 'json',
        'image_quality': 'medium',
        'ocr_language': 'chi_sim',
        'concurrent_processes': '3',
        'vector_dimension': '1536',
        'similarity_threshold': '0.8',
        'theme': 'light',
        'font_size': 'medium',
        'language': 'zh',
        'animations': 'enabled',
        'auto_save': 'enabled',
        'auto_save_interval': '5',
        'max_file_size': '50',
        'temp_file_retention': '7',
        'log_level': 'info',
        'log_retention': '30'
    }
}

# 定义默认的上传目录和允许的文件扩展名
UPLOAD_DIR = DEFAULT_CONFIG['Settings']['upload.dir']
ALLOWED_EXTENSIONS = set(DEFAULT_CONFIG['Settings']['allowed.extensions'].split(','))

# 确保上传目录存在
os.makedirs(UPLOAD_DIR, exist_ok=True)

def load_config():
    config = configparser.ConfigParser()
    config_file = os.path.join(os.path.dirname(__file__), CONFIG_FILE)
    
    try:
        if not os.path.exists(config_file):
            # 如果配置文件不存在，创建默认配置
            config['API'] = DEFAULT_CONFIG['API']
            config['OpenAI'] = {
                'api_key': 'your_openai_api_key',
                'api_base': 'https://api.openai.com/v1',
                'model': 'gpt-4',
                'temperature': '0.7',
                'max_tokens': '2000'
            }
            config['Settings'] = DEFAULT_CONFIG['Settings']
            with open(config_file, 'w') as f:
                config.write(f)
        
        # 读取配置文件
        config.read(config_file)
        
        # 返回完整的配置内容
        return {
            'content': '\n'.join([
                '[API]',
                *[f"{key} = {value}" for key, value in config['API'].items()],
                '',
                '[OpenAI]',
                *[f"{key} = {value}" for key, value in config['OpenAI'].items()],
                '',
                '[Settings]',
                *[f"{key}={value}" for key, value in config['Settings'].items()]
            ])
        }
    except Exception as e:
        print(f"加载配置失败: {str(e)}")
        return {
            'content': '\n'.join([
                '[API]',
                *[f"{key} = {value}" for key, value in DEFAULT_CONFIG['API'].items()],
                '',
                '[OpenAI]',
                *[f"{key} = {value}" for key, value in {
                    'api_key': 'your_openai_api_key',
                    'api_base': 'https://api.openai.com/v1',
                    'model': 'gpt-4',
                    'temperature': '0.7',
                    'max_tokens': '2000'
                }.items()],
                '',
                '[Settings]',
                *[f"{key}={value}" for key, value in DEFAULT_CONFIG['Settings'].items()]
            ])
        }

def save_config(config):
    config_file = os.path.join(os.path.dirname(__file__), CONFIG_FILE)
    try:
        # 读取现有配置
        config_parser = configparser.ConfigParser()
        config_parser.read(config_file)
        
        # 更新配置
        if 'API' not in config_parser:
            config_parser['API'] = {}
        if 'Settings' not in config_parser:
            config_parser['Settings'] = {}
        
        # 更新 API 配置
        for key, value in config.get('api', {}).items():
            config_parser['API'][key] = value
        
        # 更新 Settings 配置
        for key, value in config.get('settings', {}).items():
            config_parser['Settings'][key] = str(value)
        
        # 保存配置
        with open(config_file, 'w') as f:
            config_parser.write(f)
        
        print(f"配置已保存到: {config_file}")
        return True
    except Exception as e:
        print(f"保存配置失败: {str(e)}")
        return False

@app.route('/api/config', methods=['GET'])
def get_config():
    try:
        config = load_config()
        return jsonify(config)
    except Exception as e:
        return jsonify({"success": False, "message": str(e)})

@app.route('/api/config', methods=['POST'])
def update_config():
    try:
        new_config = request.json
        print("Received config update:", new_config)  # 添加调试日志
        
        # 解析 properties 格式的配置
        config_parser = configparser.ConfigParser()
        if 'content' in new_config:
            # 从 properties 格式解析
            content = new_config['content']
            lines = content.split('\n')
            current_section = None
            
            for line in lines:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                
                if line.startswith('[') and line.endswith(']'):
                    current_section = line[1:-1]
                    if current_section not in config_parser:
                        config_parser[current_section] = {}
                    continue
                
                if current_section:
                    if current_section in ['API', 'OpenAI']:
                        # API 和 OpenAI 部分使用 "key = value" 格式
                        parts = line.split(' = ', 1)
                        if len(parts) == 2:
                            key, value = parts
                            config_parser[current_section][key.strip()] = value.strip()
                    else:
                        # Settings 部分使用 "key=value" 格式
                        parts = line.split('=', 1)
                        if len(parts) == 2:
                            key, value = parts
                            config_parser[current_section][key.strip()] = value.strip()
        
        # 确保所有必需的 OpenAI 配置项都存在
        if 'OpenAI' not in config_parser:
            config_parser['OpenAI'] = {}
        
        # 设置默认值（如果不存在）
        default_openai_config = {
            'api_key': 'your_openai_api_key',
            'api_base': 'https://api.openai.com/v1',
            'model': 'gpt-4',
            'temperature': '0.7',
            'max_tokens': '2000'
        }
        
        for key, value in default_openai_config.items():
            if key not in config_parser['OpenAI']:
                config_parser['OpenAI'][key] = value
        
        # 保存配置
        config_file = os.path.join(os.path.dirname(__file__), CONFIG_FILE)
        with open(config_file, 'w') as f:
            config_parser.write(f)
        
        print(f"配置已保存到: {config_file}")  # 添加调试日志
        return jsonify({"success": True, "message": "配置已更新"})
    except Exception as e:
        print(f"保存配置失败: {str(e)}")  # 添加调试日志
        return jsonify({"success": False, "message": str(e)})


@app.route("/")
def index():
    return render_template("manage.html")

@app.route("/convert", methods=["POST"])
def convert_file():
    try:
        if 'file' not in request.files:
            return jsonify({"success": False, "message": "没有上传文件"}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({"success": False, "message": "没有选择文件"}), 400

        if not file.filename.lower().endswith('.pdf'):
            return jsonify({"success": False, "message": "只支持PDF文件"}), 400

        # 创建临时文件
        temp_dir = tempfile.mkdtemp()
        temp_input = os.path.join(temp_dir, file.filename)
        file.save(temp_input)

        # 确保输出目录存在
        os.makedirs(UPLOAD_DIR, exist_ok=True)

        def generate_progress():
            # 执行转换命令
            cmd = ["marker_single", temp_input, "--output_dir", UPLOAD_DIR]
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                universal_newlines=True
            )

            # 读取输出
            while True:
                output = process.stdout.readline()
                if output == '' and process.poll() is not None:
                    break
                if output:
                    yield f"data: {json.dumps({'message': output.strip()})}\n\n"

            # 检查是否有错误
            if process.returncode != 0:
                error = process.stderr.read()
                yield f"data: {json.dumps({'error': error.strip()})}\n\n"
            else:
                # 转换成功，发送成功消息
                yield f"data: {json.dumps({'success': True, 'message': '转换完成，正在刷新文件列表...'})}\n\n"

            # 清理临时文件
            shutil.rmtree(temp_dir)

        return Response(generate_progress(), mimetype='text/event-stream')

    except Exception as e:
        return jsonify({"success": False, "message": f"处理文件时出错: {str(e)}"}), 500

@app.route("/list_files")
def list_files():
    if not os.path.exists(UPLOAD_DIR):
        return jsonify({"files": []})
    
    # 获取所有文件，包括子目录中的文件
    all_files = []
    for root, dirs, files in os.walk(UPLOAD_DIR):
        for file in files:
            # 获取文件扩展名
            file_ext = os.path.splitext(file)[1].lower()
            # 只添加允许的文件类型
            if file_ext in ALLOWED_EXTENSIONS:
                # 获取相对路径
                rel_path = os.path.relpath(os.path.join(root, file), UPLOAD_DIR)
                # 确保路径使用正斜杠
                rel_path = rel_path.replace('\\', '/')
                all_files.append(rel_path)
    
    # 按文件名排序
    all_files.sort()
    
    return jsonify({"files": all_files})

@app.route("/download")
def download_file():
    filename = request.args.get("filename")
    if not filename:
        return "Missing filename", 400
    
    file_path = os.path.join(UPLOAD_DIR, filename)
    if not os.path.exists(file_path):
        return "File not found", 404
    
    directory = os.path.dirname(file_path)
    basename = os.path.basename(file_path)
    
    return send_from_directory(directory, basename, as_attachment=True)

@app.route("/convert_qa", methods=["POST"])
def convert_qa():
    try:
        if 'file' not in request.files:
            return jsonify({"success": False, "message": "没有上传文件"}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({"success": False, "message": "没有选择文件"}), 400

        if not file.filename.lower().endswith('.md'):
            return jsonify({"success": False, "message": "只支持Markdown文件"}), 400

        def generate_qa_progress():
            try:
                # 读取上传的文件内容
                content = file.read().decode('utf-8')
                
                qa_pattern = re.compile(
                    r'QA \d+:\n'
                    r'Q: (.*?)\n'
                    r'A: (.*?)\n'
                    r'(?=\nQA \d+:|$)',
                    re.DOTALL
                )

                matches = qa_pattern.findall(content)
                qa_list = []

                for i, (question, answer) in enumerate(matches, 1):
                    qa_list.append({
                        "id": i,
                        "question": question.strip(),
                        "answer": answer.strip()
                    })
                    
                    # 发送进度
                    yield f"data: {json.dumps({'message': f'处理第 {i} 个QA对'})}\n\n"

                # 发送完整的QA列表
                yield f"data: {json.dumps({'success': True, 'qa_list': qa_list, 'message': f'成功提取 {len(qa_list)} 个QA对'})}\n\n"

            except Exception as e:
                yield f"data: {json.dumps({'error': f'处理QA时出错: {str(e)}'})}\n\n"

        return Response(generate_qa_progress(), mimetype='text/event-stream')

    except Exception as e:
        return jsonify({"success": False, "message": f"处理文件时出错: {str(e)}"}), 500

def load_rag_api_config_dict():
    config = configparser.ConfigParser()
    config_file = os.path.join(os.path.dirname(__file__), CONFIG_FILE)
    config.read(config_file)
    api = config['API'] if 'API' in config else {}
    return {
        'base_url': api.get('base_url', ''),
        'api_key': api.get('api_key', ''),
        'dataset_id': api.get('dataset_id', ''),
        'document_id': api.get('document_id', '')
    }

@app.route("/send_to_knowledge_base", methods=["POST"])
def send_to_knowledge_base():
    try:
        data = request.json
        qa_list = data.get("qa_list", [])
        # 只用后端读取的配置
        api_config = load_rag_api_config_dict()
        print(qa_list)
        if not qa_list:
            return jsonify({"success": False, "message": "没有QA数据"}), 400
        # 检查配置
        if not all([api_config['base_url'], api_config['api_key'], api_config['dataset_id'], api_config['document_id']]):
            return jsonify({"success": False, "message": "API配置不完整，请在设置页面填写并保存 Base URL、API Key、Dataset ID、Document ID"}), 400
        api_url = f"{api_config['base_url']}/api/v1/datasets/{api_config['dataset_id']}/documents/{api_config['document_id']}/chunks"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_config['api_key']}"
        }
        results = []
        for qa in qa_list:
            qa_content = f"Q: {qa['question']} A: {qa['answer']}"
            data = {"content": qa_content}
            try:
                response = requests.post(api_url, headers=headers, json=data)
                results.append({
                    "question": qa['question'],
                    "success": response.status_code == 200,
                    "message": response.text if response.status_code != 200 else "成功"
                })
            except Exception as e:
                results.append({
                    "question": qa['question'],
                    "success": False,
                    "message": str(e)
                })
        return jsonify({
            "success": True,
            "results": results
        })
    except Exception as e:
        return jsonify({"success": False, "message": f"发送到知识库时出错: {str(e)}"}), 500

@app.route("/preview")
def preview_file():
    filename = request.args.get("filename")
    if not filename:
        return "Missing filename", 400
    
    file_path = os.path.join(UPLOAD_DIR, filename)
    if not os.path.exists(file_path):
        return "File not found", 404
    
    # 获取文件所在目录
    directory = os.path.dirname(file_path)
    # 获取文件名
    basename = os.path.basename(file_path)
    
    # 根据文件类型返回不同的预览方式
    if filename.lower().endswith('.md'):
        # 对于 Markdown 文件，返回文本内容
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            return jsonify({
                "type": "markdown",
                "content": content
            })
        except Exception as e:
            return jsonify({"error": str(e)}), 500
    elif filename.lower().endswith('.pdf'):
        # 对于 PDF 文件，返回文件 URL
        return jsonify({
            "type": "pdf",
            "url": f"/download?filename={filename}"
        })
    else:
        # 对于其他文件类型，返回文件 URL
        return jsonify({
            "type": "other",
            "url": f"/download?filename={filename}"
        })

@app.route("/extract_qa", methods=["POST"])
def extract_qa():
    try:
        data = request.json
        content = data.get("content")
        prompt = data.get("prompt")
        
        if not content or not prompt:
            return jsonify({"success": False, "message": "缺少必要参数"}), 400

        def generate_qa_progress():
            try:
                # 读取配置
                config = configparser.ConfigParser()
                config.read(os.path.join(os.path.dirname(__file__), CONFIG_FILE))
                
                # 构建提示词
                system_prompt = """你是一个专业的QA提取助手。请根据给定的规则从文本中提取QA对。
                规则说明：
                1. 仔细分析文本内容
                2. 提取有意义的问题和答案对
                3. 确保问题和答案的完整性
                4. 保持专业性和准确性
                5. 保持原文的语言，不要翻译。

                生成示例：
                Q: 你是谁？
                A: 我是一个AI助手，专门用于从文本中提取QA对。
                """

                # 调用 OpenAI API
                openai.api_key = config['OpenAI']['api_key']
                openai.api_base = config['OpenAI']['api_base']
                response = openai.ChatCompletion.create(
                    model=config['OpenAI']['model'],
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": f"提取规则：\n{prompt}\n\n文本内容：\n{content}"}
                    ],
                    temperature=float(config['OpenAI']['temperature']),
                    max_tokens=int(config['OpenAI']['max_tokens']),
                    stream=True
                )

                # 直接输出模型响应
                for chunk in response:
                    print(chunk)
                    # finish_reason 是 stop 停止
                    if chunk.choices[0].finish_reason == 'stop':
                        break
                    if chunk.choices[0].delta.content:
                        message = chunk.choices[0].delta.content
                        yield f"data: {json.dumps({'message': message})}\n\n"

            except Exception as e:
                yield f"data: {json.dumps({'error': f'提取QA时出错: {str(e)}'})}\n\n"

        return Response(generate_qa_progress(), mimetype='text/event-stream')

    except Exception as e:
        return jsonify({"success": False, "message": f"处理请求时出错: {str(e)}"}), 500

if __name__ == "__main__":
    os.makedirs(FILES_DIR, exist_ok=True)
    app.run(host="0.0.0.0", port=5001, debug=True)